//
// Created by user on 20-3-20.
//

#include "Task.h"
#include "SMP.h"
#include "Schedule.h"
//#include "../MemoryManagement/Malloc.h"

union TaskUnion initTaskUnion __attribute__((__section__ ("_initTaskUnion"))) = {INIT_TASK(initTaskUnion.task)};

MemorySpaceDistribution memorySpaceDistribution = {0};

// 根据CPU的核心数量，来初始化线程结构体的数量
Task *initTask[NR_CPUS] = {&initTaskUnion.task, 0};

Thread initThread =
        {
                .rsp0 = (unsigned long) (initTaskUnion.stack + STACK_SIZE / sizeof(unsigned long)),
                .rsp = (unsigned long) (initTaskUnion.stack + STACK_SIZE / sizeof(unsigned long)),
                .fs = KERNEL_DS,
                .gs = KERNEL_DS,
                .cr2 = 0,
                .trapNumber = 0,
                .errorCode = 0
        };

Tss initTss[NR_CPUS] = {[0 ... NR_CPUS - 1] = INIT_TSS};


inline Task *getCurrent() {

    Task *currentTask = NULL;

    // -32767 = 32KB - 1
    __asm__ __volatile__ ("andq %%rsp,%0	\n\t":"=r"(currentTask):"0"(~32767UL));

    return currentTask;
}


// 切换任务状态段和FS、GS寄存器
void __switch_to(Task *prevTask, Task *nextTask) {

    unsigned int color = 0;

    // 1. 把任务状态段的第1个的 rsp0 设置为下一个线程的 rsp0
    initTss[SMPCPUID()].rsp0 = nextTask->thread->rsp0;


    // 2. 然后把任务状态段的第1个的信息设置到 CPU 中
//    setTSS64(initTss[0].rsp0, initTss[0].rsp1, initTss[0].rsp2, initTss[0].ist1, initTss[0].ist2, initTss[0].ist3,
//             initTss[0].ist4, initTss[0].ist5, initTss[0].ist6, initTss[0].ist7);

    // 3. 将当前的 GS 和 FS 寄存器的值保存到上个线程的 FS 和 GS 位
    __asm__ __volatile__("movq	%%fs,	%0 \n\t":"=a"(prevTask->thread->fs));
    __asm__ __volatile__("movq	%%gs,	%0 \n\t":"=a"(prevTask->thread->gs));

    // 4. 将下一个线程的 FS 和 GS 设置到 FS 和 GS 寄存器
    __asm__ __volatile__("movq	%0,	%%fs \n\t"::"a"(nextTask->thread->fs));
    __asm__ __volatile__("movq	%0,	%%gs \n\t"::"a"(nextTask->thread->gs));

    wrmsr(0x175, nextTask->thread->rsp0);

    if(SMPCPUID() == 0)
        color = WHITE;
    else
        color = YELLOW;

    PerformSite *regs = (PerformSite *) ((unsigned long) nextTask + STACK_SIZE - sizeof(PerformSite));
//    printfColor(color, BLACK, "regs->rip: %lx\n", regs->rip);

    printfColor(color, BLACK, "prevTask->thread->rsp0: %lx\n", prevTask->thread->rsp0);
    printfColor(color, BLACK, "nextTask->thread->rsp0: %lx\n", nextTask->thread->rsp0);

    printfColor(color, BLACK, "prevTask: %lx\n", prevTask);
    printfColor(color, BLACK, "nextTask: %lx\n", nextTask);
//    printfColor(color, BLACK, "prevTask->thread->rip: %lx\n", prevTask->thread->rip);
//    printfColor(color, BLACK, "nextTask->thread->rip: %lx\n", nextTask->thread->rip);
    printfColor(color, BLACK, "nextTask->thread: %lx\n", nextTask->thread);


}

void taskInit() {

    Task *pTask = NULL;

    // 1. 先设置全局变量 内存空间分布 的值，主要是设置内核文件在内存中的空间分布情况
    memorySpaceDistribution.pgd = (pml4t_t *) globalCR3;

    memorySpaceDistribution.startCode = memoryManagementStruct.kernelStartCode;
    memorySpaceDistribution.endCode = memoryManagementStruct.kernelEndCode;

    memorySpaceDistribution.startData = (unsigned long) &_data;
    memorySpaceDistribution.endData = memoryManagementStruct.kernelEndData;

    memorySpaceDistribution.startRodata = (unsigned long) &_rodata;
    memorySpaceDistribution.endRodata = (unsigned long) &_erodata;

    memorySpaceDistribution.startBrk = 0;
    memorySpaceDistribution.endBrk = memoryManagementStruct.kernelEndBrk;

    memorySpaceDistribution.startStack = _stack_start;

    // 把内核代码段写入到 MSR 特殊寄存器里面
    wrmsr(0x174, KERNEL_CS);
    wrmsr(0x175, current->thread->rsp0);
    wrmsr(0x176, (unsigned long) systemCall);

    // 2. 设置进程的任务状态段

    // 3. 根据处理器的ID设置默认的 rsp0
    initTss[SMPCPUID()].rsp0 = initThread.rsp0;

    // 4. 初始化内核栈底的进程结构体
    listInit(&initTaskUnion.task.list);

    // 5. 创建进程的最后最后一步，设置指定的寄存器和分配内存
    kernelThread(init, 10, CLONE_FS | CLONE_FILES | CLONE_SIGNAL);

    // 6. 设置进程的状态位运行状态
    initTaskUnion.task.preemptCount = 0;
    initTaskUnion.task.state = TASK_RUNNING;

    // 7. 获取下一个进程的 TASK 结构体指针
//    pTask = containerOf(listNext(&taskSchedule[0].taskQueueHead.list), Task, list);
//
//    printfColor(INDIGO, BLACK, "pTask->thread->rsp = %lx\n", pTask->thread->rsp);
//
//    // 8. 切换进程
//    switchTo(current, pTask);

}

unsigned long systemCallFunction(PerformSite *regs) {

    printfColor(INDIGO, BLACK, "systemCallFunction\n");

    return systemCallTable[regs->rax](regs);

}

void userLevelFunction() {

    long ret = 0;
    char string[] = "Hello World!\n";

//    printfColor(BLUE, BLACK, "1\n");

    // 执行 sysexit 指令需要将返回地址和栈顶地址分别放到 rdx 和 rcx 中
    __asm__ __volatile__    (
    "leaq	sysexitReturnAddress(%%rip),    %%rdx	\n\t"
            "movq	%%rsp,	%%rcx		\n\t"
            "sysenter			\n\t"
            "sysexitReturnAddress:	\n\t"
    :"=a"(ret):"0"(1), "D"(string):"memory");

    for (int i = 0; i < 10; ++i) {
        i = i + 1;
    }

    while (1);

}

unsigned long doExecve(PerformSite *regs) {

    unsigned long addr = 0x800000;
    unsigned long rspAddr = 0xa00000;
    unsigned long *tmp;
    unsigned long *virtual = NULL;
    Page *p = NULL;

    regs->rdx = addr;    // RIP
    regs->rcx = 0xa00000;    // RSP
    regs->rax = 1;
    regs->ds = 0;
    regs->es = 0;

//    printfColor(INDIGO, BLACK, "regs: %lx\n", regs);
    printfColor(RED, BLACK, "doExecve task is running\n");

    globalCR3 = getCR3();

    // 为代码段分配物理页
    tmp = Physical_To_Virtual(
            (unsigned long *) ((unsigned long) globalCR3 & (~0xfffUL)) + ((addr >> PAGE_GDT_SHIFT) & 0x1ff));

    virtual = kmalloc(PAGE_4K_SIZE, 0);
    set_mpl4t(tmp, mk_mpl4t(Virtual_To_Physical(virtual), PAGE_USER_GDT));

    tmp = Physical_To_Virtual((unsigned long *) (*tmp & (~0xfffUL)) + ((addr >> PAGE_1G_SHIFT) & 0x1ff));
    virtual = kmalloc(PAGE_4K_SIZE, 0);
    set_pdpt(tmp, mk_pdpt(Virtual_To_Physical(virtual), PAGE_USER_Dir));

    tmp = Physical_To_Virtual((unsigned long *) (*tmp & (~0xfffUL)) + ((addr >> PAGE_2M_SHIFT) & 0x1ff));
    p = allocPages(ZONE_NORMAL, 1, PG_PTable_Maped);
    set_pdt(tmp, mk_pdt(p->physicalAddress, PAGE_USER_Page));

    // 为栈分配物理页
//    tmp = Physical_To_Virtual(
//            (unsigned long *) ((unsigned long) globalCR3 & (~0xfffUL)) + ((rspAddr >> PAGE_GDT_SHIFT) & 0x1ff));
//
//    virtual = kmalloc(PAGE_4K_SIZE, 0);
//    set_mpl4t(tmp, mk_mpl4t(Virtual_To_Physical(virtual), PAGE_USER_GDT));
//
//    tmp = Physical_To_Virtual((unsigned long *) (*tmp & (~0xfffUL)) + ((rspAddr >> PAGE_1G_SHIFT) & 0x1ff));
//    virtual = kmalloc(PAGE_4K_SIZE, 0);
//    set_pdpt(tmp, mk_pdpt(Virtual_To_Physical(virtual), PAGE_USER_Dir));
//
//    tmp = Physical_To_Virtual((unsigned long *) (*tmp & (~0xfffUL)) + ((rspAddr >> PAGE_2M_SHIFT) & 0x1ff));
//    p = allocPages(ZONE_NORMAL, 1, PG_PTable_Maped);
//    set_pdt(tmp, mk_pdt(p->physicalAddress, PAGE_USER_Page));


    flushTLB();

    // 限制进程访问内核的内存
    if (!(current->flags & PF_KTHREAD))
        current->addressLimit = 0xffff800000000000;

    memcpy((void *) addr, userLevelFunction, 1024);

//    printfColor(INDIGO, BLACK, "regs: %lx\n", regs->rdx);

    return 1;
}

unsigned long init(unsigned long arg) {

    PerformSite *regs;
    // 提示一下进程已经运行了
    printfColor(RED, BLACK, "init task is running, arg: %x\n", arg);

    current->thread->rip = (unsigned long) returnSystemCall;
    current->thread->rsp = (unsigned long) current + STACK_SIZE - sizeof(PerformSite);
    regs = (PerformSite *) current->thread->rsp;

    current->flags = 0;

    printfColor(INDIGO, BLACK, "regs: %lx\n", regs);

    // 通过向栈顶
    __asm__ __volatile__    (    "movq	%1,	%%rsp	\n\t"
            "movq	%0,	%%rdi	\n\t"
            "pushq	%2		\n\t"
            "jmp	doExecve	\n\t"
    ::"D"(regs), "m"(current->thread->rsp), "m"(current->thread->rip):"memory");

    return 1;
}

unsigned long kernelThread(unsigned long (*function)(unsigned long), unsigned long arg, unsigned long flags) {

    PerformSite regs;

    // 1. 清空 regs
    memset(&regs, 0, sizeof(regs));

    // 2. 设置 rbx 放进程运行的第一个函数
    regs.rbx = (unsigned long) function;

    // 3. 设置 rdx 放传给进程第一个函数的参数
    regs.rdx = arg;

    // 4. 设置进程的段寄存器
    regs.ds = KERNEL_DS;
    regs.es = KERNEL_DS;
    regs.cs = KERNEL_CS;
    regs.ss = KERNEL_DS;

    printfColor(WHITE, BLACK, "function = %lx\n", function);

    // 5. 设置进程的标志寄存器
    regs.rflags = (1 << 9);

    // 6. 设置进程默认的执行入口，这个是内核层的执行入口
    regs.rip = (unsigned long) kernelThreadFunc;

    return doFork(&regs, flags, 0, 0);
}

unsigned long doFork(PerformSite *regs, unsigned long cloneFlags, unsigned long stackStart, unsigned long stackSize) {

    Task *tsk = NULL;
    Thread *thd = NULL;
    Page *p = NULL;

    // 1. 为进程申请一个物理页
    p = allocPages(ZONE_NORMAL, 1, PG_PTable_Maped | PG_Active | PG_Kernel);

    // 2. 将刚申请到的物理地址映射为内核的虚拟地址
    tsk = (Task *) Physical_To_Virtual(p->physicalAddress);

//    printfColor(WHITE, BLACK, "Task pid = %x\n", tsk->pid);

    // 3. 将新申请到的物理页全部置为0
    memset(tsk, 0, sizeof(*tsk));

    // 4. 将当前 TASK 的内容设置到刚申请的物理页里面
    *tsk = *current;

    // 5. 初始化刚申请的 tsk 的链表
    listInit(&tsk->list);

    // 6. 将 tsk 的链表加入到进程队列里面
    listAddToBefore(&initTaskUnion.task.list, &tsk->list);

    tsk->priority = 2;
    tsk->cpuID = SMPCPUID();
    // 7. 将进程ID号加一
    tsk->pid++;

    // 8. 设置进程状态为不间断的
    tsk->state = TASK_UNINTERRUPTIBLE;

    // 9. tsk + 1 会得到 tsk 所占据的内存的首地址，然后把这个首地址当做线程结构体的首地址
    thd = (Thread *) (tsk + 1);
    memset(thd, 0, sizeof(*thd));
    tsk->thread = thd;

    // 10. 将在 kernelThread 函数中设置好的 regs 复制到指定位置
    memcpy((void *) ((unsigned long) tsk + STACK_SIZE - sizeof(PerformSite)), regs, sizeof(PerformSite));
//    memcpy(regs, (void *) ((unsigned long) tsk + STACK_SIZE - sizeof(PerformSite)), sizeof(PerformSite));

    // 11. 设置进程的线程信息
    thd->rsp0 = (unsigned long) tsk + STACK_SIZE;
    thd->rip = regs->rip;
    thd->rsp = (unsigned long) tsk + STACK_SIZE - sizeof(PerformSite);
    thd->fs = 0;
    thd->gs = 0;

    // 12. 判断进程是否运行在应用层层，如果是就到 returnSystemCall 执行
    if (!(tsk->flags & PF_KTHREAD)) {
        thd->rip = regs->rip = (unsigned long) returnSystemCall;
    }

    // 13. 更改进程的状态为运行状态
    tsk->state = TASK_RUNNING;

    insertTaskQueue(tsk);

    return 0;
}

unsigned long doExit(unsigned long code) {

    printfColor(RED, BLACK, "exit task is running, arg: %lx\n", code);

    while (1);

}